import os
import base64
import uuid
from io import BytesIO
from PIL import Image
from typing import Optional, Tuple, Union
import mimetypes

class ImageProcessor:
    """图片处理器"""
    
    def __init__(self, upload_folder: str, allowed_extensions: set):
        """
        初始化图片处理器
        
        Args:
            upload_folder: 上传文件夹路径
            allowed_extensions: 允许的文件扩展名
        """
        self.upload_folder = upload_folder
        self.allowed_extensions = allowed_extensions
        
        # 确保上传目录存在
        if not os.path.exists(upload_folder):
            os.makedirs(upload_folder)
    
    def is_allowed_file(self, filename: str) -> bool:
        """检查文件扩展名是否允许"""
        return '.' in filename and \
               filename.rsplit('.', 1)[1].lower() in self.allowed_extensions
    
    def validate_image(self, image_data: Union[str, bytes]) -> Tuple[bool, str]:
        """
        验证图片数据
        
        Args:
            image_data: 图片数据（base64字符串或字节）
            
        Returns:
            (是否有效, 错误信息)
        """
        try:
            if isinstance(image_data, str):
                # 处理base64数据
                if image_data.startswith('data:image/'):
                    # 移除data URL前缀
                    image_data = image_data.split(',', 1)[1]
                
                # 解码base64
                image_bytes = base64.b64decode(image_data)
            else:
                image_bytes = image_data
            
            # 尝试打开图片
            image = Image.open(BytesIO(image_bytes))
            image.verify()  # 验证图片完整性
            
            # 检查图片尺寸
            if image.size[0] > 4096 or image.size[1] > 4096:
                return False, "图片尺寸过大，最大支持4096x4096像素"
            
            # 检查文件大小（16MB限制）
            if len(image_bytes) > 16 * 1024 * 1024:
                return False, "图片文件过大，最大支持16MB"
            
            return True, ""
            
        except Exception as e:
            return False, f"图片格式无效: {str(e)}"
    
    def process_base64_image(self, base64_data: str) -> Tuple[Optional[str], Optional[str]]:
        """
        处理base64图片数据
        
        Args:
            base64_data: base64编码的图片数据
            
        Returns:
            (保存的文件路径, 错误信息)
        """
        try:
            # 验证图片
            is_valid, error_msg = self.validate_image(base64_data)
            if not is_valid:
                return None, error_msg
            
            # 处理data URL
            if base64_data.startswith('data:image/'):
                header, base64_data = base64_data.split(',', 1)
                # 从header中提取文件类型
                mime_type = header.split(';')[0].split(':')[1]
                extension = mimetypes.guess_extension(mime_type)
                if not extension:
                    extension = '.png'
            else:
                extension = '.png'  # 默认扩展名
            
            # 解码图片
            image_bytes = base64.b64decode(base64_data)
            
            # 生成唯一文件名
            filename = f"{uuid.uuid4()}{extension}"
            filepath = os.path.join(self.upload_folder, filename)
            
            # 保存文件
            with open(filepath, 'wb') as f:
                f.write(image_bytes)
            
            return filepath, None
            
        except Exception as e:
            return None, f"处理图片失败: {str(e)}"
    
    def process_uploaded_file(self, file) -> Tuple[Optional[str], Optional[str]]:
        """
        处理上传的文件
        
        Args:
            file: Flask上传的文件对象
            
        Returns:
            (保存的文件路径, 错误信息)
        """
        try:
            if not file or file.filename == '':
                return None, "未选择文件"
            
            if not self.is_allowed_file(file.filename):
                return None, f"不支持的文件类型，支持的格式: {', '.join(self.allowed_extensions)}"
            
            # 读取文件数据
            file_data = file.read()
            
            # 验证图片
            is_valid, error_msg = self.validate_image(file_data)
            if not is_valid:
                return None, error_msg
            
            # 生成唯一文件名
            extension = '.' + file.filename.rsplit('.', 1)[1].lower()
            filename = f"{uuid.uuid4()}{extension}"
            filepath = os.path.join(self.upload_folder, filename)
            
            # 保存文件
            with open(filepath, 'wb') as f:
                f.write(file_data)
            
            return filepath, None
            
        except Exception as e:
            return None, f"处理上传文件失败: {str(e)}"
    
    def image_to_base64(self, image_path: str) -> Optional[str]:
        """
        将图片文件转换为base64字符串
        
        Args:
            image_path: 图片文件路径
            
        Returns:
            base64编码的图片数据
        """
        try:
            with open(image_path, 'rb') as f:
                image_data = f.read()
            
            # 获取MIME类型
            mime_type, _ = mimetypes.guess_type(image_path)
            if not mime_type:
                mime_type = 'image/png'
            
            # 编码为base64
            base64_data = base64.b64encode(image_data).decode('utf-8')
            
            # 返回data URL格式
            return f"data:{mime_type};base64,{base64_data}"
            
        except Exception as e:
            print(f"转换图片为base64失败: {e}")
            return None
    
    def cleanup_file(self, filepath: str) -> bool:
        """
        清理临时文件
        
        Args:
            filepath: 文件路径
            
        Returns:
            是否成功删除
        """
        try:
            if os.path.exists(filepath):
                os.remove(filepath)
                return True
            return False
        except Exception as e:
            print(f"删除文件失败: {e}")
            return False
